/**
  ETFAna project, Anyang Normal University && IMP-CAS
  \class ETFPID
  \brief class for particle identification. This class takes tracks on both sides
  of the dipole magnet to extract particle A/Z, and particle velocity beta.
  \author SUN Yazhou, asia.rabbit@163.com
  \since 2021-11-10
  \date 2021-12-15 last modified
  \attention
  changelog
  <table>
  <tr>  <th>Date         <th>Author      <th>Description                   </tr>
  <tr>  <td>2021-11-10   <td>Asia Sun    <td>file created                  </tr>
  </table>

  \copyright Copyright (c) 2021-2024 Anyang Normal U. && IMP-CAS with LGPLv3 LICENSE
*/

#include <iostream>
#include <TGraph.h>
#include <TLegend.h>
#include <TRandom3.h>
#include <TClonesArray.h>
#include <TTree.h>
#include "ETFRun.h"
#include "ETFPID.h"
#include "ETFTrackManager.h"
#include "ETFDetectorManager.h"
#include "ETFTaTracker.h"
#include "ETFPostMagTracker.h"
#include "ETFGlobalTracker.h"
#include "ETFTrackTa4.h"
#include "ETFTrackTa.h"
#include "ETFTrack.h"
#include "ETF3DTrack.h"
#include "ETFTOFWall.h"
#include "ETFTimer2.h"
#include "ETFTimer1.h"
#include "ETFMath.h"
#include "ETFMsg.h"
#include "ETFPlaHit.h"
#include "ETFParaManager.h"
#include "ETFHPTDCCh.h"

#define sum2 ETFMath::Sum2

using std::cout;
using std::endl;

ETFPID::ETFPID() : ETFVPID(), fIsParaMom(false){}

ETFPID::~ETFPID(){}

void ETFPID::Initialize(){
  ETFVPID::Initialize();

  ftof1 = -999.;
  fBG0 = fGamma0 = fBeta0 = -1.;

  fAoZ = -999.; fBrho = -999.;
  fZO = -999.; fXO = -999.;
  fRho = -999.;
  fTrkLenT = -999.; ftof2 = -999.;
  fBG = fGamma = fBeta = -1.;
  fPoZ = -999.;

  // x2: abscissa of Mag exit point from real trk and calculated arc
  // as a PID estimator; 0-1: real-calcu
  fX2Arr[0] = -999.; fX2Arr[1] = -999.;

  fCospl = fdCospl = fPoZpl = fdPoZpl = fBGC = fdBGC = -999.;
} // end member function Initialize

void ETFPID::Configure(){
  ETFVPID::Configure();
  if((*fpm->GetMisPara())["isParaM"])
    fIsParaMom = (*fpm->GetMisPara())["isParaM"].as<bool>();
} // end member function Configure

// register to the main tree
void ETFPID::Register(){
  fTree->Branch("pid", this);
} // end member function Register

// the same fly, but collecting information using singleton classes in ETFAna
// istof: whether we only call fly to obtain tof
bool ETFPID::Fly(bool istof){
  if(fHasFlied) return fFlySuccess; // has called

  // gt is preferred to ptt
  static ETFGlobalTracker *gt = ftm->GetGlobalTracker();
  static ETFTaTracker *ptt = ftm->GetPostTaTracker();
  static ETFPostMagTracker *pmt = ftm->GetPostMagTracker();

  // even tof2 cannot be calculated without (only 1) trk-postMag
  if(!pmt || 1 != pmt->GetNofTracks()){ fHasFlied = true; return false; }
  // then we can calculate tof2 //
  ETFTrack *trkpm = dynamic_cast<ETFTrack *>(pmt->GetTrack(0));
  const int stripId = trkpm->GetFiredStripId();
  // in case that some primary beam exps don't have a tstart //
  ETFTimer *tstart = nullptr;
  if(fdm->HasDetector("tstart"))
    tstart = dynamic_cast<ETFTimer1 *>(fdm->GetDetector("tstart"));
  else if(fdm->HasDetector("tstart2"))
    tstart = dynamic_cast<ETFTimer2 *>(fdm->GetDetector("tstart2", true));
  ETFTimer2 *tstop = dynamic_cast<ETFTimer2 *>(fdm->GetDetector("tstop"));
  ETFTOFWall *tofw = nullptr; // tofw has 2 versions: tofw and tofwB
  if(fdm->HasDetector("tofw"))
    tofw = dynamic_cast<ETFTOFWall *>(fdm->GetDetector("tofw"));
  else tofw = dynamic_cast<ETFTOFWall *>(fdm->GetDetector("tofwB", true));
  // not properly fired -- not even ready to give TOF information
  double t0 = -999., t1 = -999.; // time of tstart and tstop
  if(dynamic_cast<ETFTimer2 *>(tstart)){ // both-ended timer
    if(4 == tstart->GetFiredStatus()) t0 = tstart->GetTime();
  }
  else if(tstart && tstart->GetFiredStatus())  t0 = tstart->GetTime();
  if(tstop && 4 == tstop->GetFiredStatus()) t1 = tstop->GetTime();
  if(ftof1 < 0.) if(t0 > 0. && t1 > 0.) ftof1 = t1 - t0;
  if(t1 > 0.) ftof2 = tofw->GetStrip(stripId)->GetTime() - t1 + ftof2_ofst;
  if(ftof1 > 0.) fBeta0 = L / (ftof1*c0);
  if(fBeta0 >= 1. || fBeta0 < 0.){ // abnormal beta0
    // fBeta0 = -1.;
  } // end if
  else{
    fBG0 = ETFMath::BetaGamma(fBeta0);
    fGamma0 = fBG0/fBeta0;
  } // end else
  if(istof) return true;

  // PID condition: only 1 track in globalTracker and only 1 track in postMagTracker
  // const int npreTaTrk = ftm->GetPreTaTracker()->GetNofTracks(); // DEBUG
  // if(npreTaTrk != 1){ fHasFlied = true; return false; } // DEBUG
  ETFVTrack *trkpt = nullptr; // postTaTrk
  if(gt){
    if(1 == gt->GetNofTracks()){
      ETFTrackTa4 * trkgt = dynamic_cast<ETFTrackTa4 *>(gt->GetTrack(0));
      trkpt = trkgt->GetPostTaTrack();
    } // end if
    else{ fHasFlied = true; return false; }
  } // end if
  else if(ptt && 1 == ptt->GetNofTracks()) trkpt = ptt->GetTrack(0);
  else{ fHasFlied = true; return false; }
  const double p0[2] = {trkpt->GetK(), trkpt->GetB()};
  const double p1[2] = {trkpm->GetK(), trkpm->GetB()};

  fFlySuccess = Fly(p0, p1, ftof2);
  if(fIsParaMom && fFlySuccess) ParallelMomentum(); // calculate the parallel momentum of the residue

  // for those calls only calculating tof, a 2nd chance of calling Fly is reserved
  return fFlySuccess;
} // end member function Fly

// propogate the particle in the dipole magnet to meet the tracks given by the DCs
// currently and usually a uniform magnetic field is assumed, as it is both
// appropriate and efficient
// tof2: from the target to TOFWall
// p0[0-1]: {k, b} from postTaTrk
// p1[0-1]: {k, b} from postMagTrk
bool ETFPID::Fly(const double *p0, const double *p1, double tof2){
  if(fHasFlied) return true;

  if(tof2 < 0.) return false;
  // rough estimate particle beta to decide whether it is appropriate to proceed PID //
  const double xe = p1[0]*kzMagOut+p1[1]; // x of the point exiting the magField
  double pe[3]{}; // the position of TOFWall hit
  ftm->GetPostMagTracker()->GetTOFWall()->GetHitPosition(p1, pe);
  // do the particle propogation in magField assuming a uniform magField //
  fRho = ETFMath::rho(p0[0], p0[1], p1[0], p1[1], fX2Arr, &fZO, &fXO); // in mm
  fTrkLenT = (kzMagIn-fTstopZ)*sqrt(1.+p0[0]*p0[0]) + fabs(fRho*(atan(p1[0])-atan(p0[0])))
    + sqrt(sum2(kzMagOut-pe[2], xe-pe[0]));
  fBeta = fTrkLenT / (c0*tof2);
  if(fBeta >= 1. || fBeta < 0.){ // abnormal beta
    // fBeta = -1.;
    return false;
  } // end if
  fBG = ETFMath::BetaGamma(fBeta); fGamma = fBG/fBeta;
  if(-999. == kB) ETFMsg::Error("ETFPID", "Fly: kB not assigned");
  fBrho = kB*fRho*1.e-3; // in T.m; r: 1e-3 factor for mm->m
  fPoZ = fBrho*c0; // in MeV/c
  // CC = 0.321840605 = e0/(u0*c0*1E6) SI unit
  fAoZ = fBrho*CC/fBG;

  fHasFlied = true;
  return true;
} // end member function Fly

// for calibration purposes
double ETFPID::poz1() const{
  if(fRho < 0.) return -999.;
  const double brho = 1.074*fRho*1.e-3; // in T.m
  const double poz = brho*c0; // in MeV/c
  return poz; //  
} // end member function poz1

double ETFPID::beta1() const{
  // if(ftof2 < 0. || fTrkLenT < 0.) return -999.;
  // const double beta = fTrkLenT / (c0*(ftof2+1.));
  const double bg0 = bg1();
  if(bg0 <= 0.) return -999.;
  const double beta = bg0 / sqrt(1.+bg0*bg0);
  if(beta >= 1. || beta <= 0.) return -999.;
  return beta;
} // end member function beta1

double ETFPID::bg1() const{
  // const double beta = beta1();
  // if(beta >= 1. || beta <= 0. || fRho < 0.) return -999.;
  // const double bg = ETFMath::BetaGamma(beta);
  const double bg0 = bg();
  if(bg0 <= 0.) return -999.;
  // bg from measured (postTA) to calibrated (preTA)
  const double bg = bg0;
  return bg;
} // end member function bg1

double ETFPID::aoz1() const{
  // const double poz = poz1(), bg = bg1();
  // if(poz < 0. || bg <= 0.) return -999.;
  // const double brho = poz / c0;
  // const double aoz = brho * CC / bg;
  if(-999. == fAoZ) return -999.;
  // return fAoZ * 0.925459 + 0.0718586;
  return fAoZ * 0.738185 + 0.494146;
} // end member function aoz1

// x2arr[0] - x2arr[1]
double ETFPID::dx2() const{
  if(-999. == fX2Arr[0] || -999. == fX2Arr[1]) return -999.;
  return fX2Arr[0] - fX2Arr[1];
} // end member function dx2

// calculate the parallel momentum of the residue
// cos(theta_core - theta_projectile)
bool ETFPID::ParallelMomentum(){
  static ETFTaTracker *preTaTracker = ftm->GetPreTaTracker();
  static ETFTaTracker *postTaTracker = ftm->GetPostTaTracker();

  // validity check
  if(1 != preTaTracker->GetNof3DTracks()) return false;
  if(1 != postTaTracker->GetNof3DTracks()) return false;
  if(-999. == fPoZ) return false;

  // get the direction vector of the projectile: n0
  // and the reaction residue (the core): n1
  ETF3DTrack *trk0 = dynamic_cast<ETF3DTrack *>(preTaTracker->Get3DTrack(0));
  ETF3DTrack *trk1 = dynamic_cast<ETF3DTrack *>(postTaTracker->Get3DTrack(0));
  const double k01 = trk0->GetK1(), k02 = trk0->GetK2();
  const double k11 = trk1->GetK1(), k12 = trk1->GetK2();
  const double n0[3] = {1, k01, k02}, n1[3] = {1, k11, k12};
  const double n0m = ETFMath::norm(n0), n1m = ETFMath::norm(n1); // |n0| and |n1|
  // calculate cos<n0, n1> in fCospl        pl: parallel
  fCospl = ETFMath::innerProd(n0,n1) / (n0m*n1m); // n0*n1/(|n0|*|n1|)=cos<n0,n1>
  // calculate dcospl to see its range of magnitude //
  // slope resolution (sigma) of core (C) track and projectile (P) track around TA
  static const double DKC = 0.0001, DKP = 0.0001; // preTa: 3.07201e-03
  const double t1 = DKP*n1m*n1m, t2 = DKC*n0m*n0m;
  fdCospl = sqrt(sum2(
    // dcospl/dk01   dcospl/dk02
    // dcospl/dk11   dcospl/dk12
    t1*((1.+k02*k02)*k11-k01*(1.+k02*k12)),t1*((1.+k01*k01)*k12-k02*(1.+k01*k11)),
    t2*((1.+k12*k12)*k01-k11*(1.+k12*k02)),t2*((1.+k11*k11)*k02-k12*(1.+k11*k01))));
  fdCospl /= pow(n0m*n1m,3);
  // calculate dpozpl to see its range of magnitude //
  static const double DBG = 2.00767e-3; // sigma(beta*gamma)
  /// fdBGC*aoz=fdpozpl, yet more accurate (aoz resolution excluded)
  /// by using real aoz after identifying the nuclide
  fBGC = fBG*fCospl*u0; // fBGC*aoz gives more accurate fPoZpl with real aoz
  fdBGC = fBGC*sqrt(sum2(DBG/fBG, fdCospl/fCospl));
  static const double DRHOREL = 0.01; // sigma(rho)/rho ~ 0.01 for ETF
  fPoZpl = fPoZ * fCospl; // projection to the incident direction
  fdPoZpl = fPoZpl*sqrt(sum2(fdCospl/fCospl, DRHOREL));
  // cout << endl << "DBG/fBG: " << DBG/fBG; // DEBUG
  // cout << " fdCospl/fCospl: " << fdCospl/fCospl << endl; // DEBUG
  // cout << "(DBG/fBG)/(fdCospl/fCospl): " << (DBG/fBG)/(fdCospl/fCospl) << endl; // DEBUG
  // cout << "fdBGC/fBGC: " << fdBGC/fBGC << endl; // DEBUG
  // cout << "fdPoZpl/fPoZpl: " << fdPoZpl/fPoZpl << endl; // DEBUG
  // getchar(); // DEBUG

  return true;
} // end member function ParallelMomentum

void ETFPID::ConstructGraph(){
  ETFVPID::ConstructGraph();

  static const int NPT = 400; // nof points to draw
  static TRandom3 rdm; ///< for generating random points

  // label the aoz value in the graph
  if(!fHasFlied || -999. == fAoZ) return; // no eligible for drawing
  if(!fGCurve) ETFMsg::Error("ETFPID", "ConstructGraph: fGCurve is nullptr");

  const int np = fGCurve->GetN(); // # points already in fGCurve
  int n = np;
  while(n - np <= NPT){
    const double ai = rdm.Uniform(-HPI, 0.); // (-pi/2, 0)
    const double zi = fZO + fRho*cos(ai);
    const double xi = fXO + fRho*sin(ai);
    if(zi < kzMagIn || zi > kzMagOut) continue; // only draw the arc
    fGCurve->SetPoint(n++, zi, xi);
  } // end for over i
  char name0[64]; sprintf(name0, "aoz_%1.3f\n", fAoZ);
  fGLegend->AddEntry(fGCurve, name0, "lp");
} // end member function ConstructGraph


void ETFPID::Print(){
  ETFVPID::Print();
  cout << "\033[32;1m============= PID PRINT (Event " << Index() << ") =============\033[0m" << endl;
  if(-999. == fAoZ){
    cout << "No particle is identified." << endl;
    return;
  } // end if

  cout << "aoz: " << fAoZ << ", beta: " << fBeta;
  cout << ", poz: " << fPoZ << " MeV/c, fBrho: " << fBrho << " Tm" << endl;
  cout << "total track length (during tof2): " << fTrkLenT << " mm" << endl;
  cout << "tof1: " << ftof1 << " ns,  tof2: " << ftof2 << " ns" << endl;
  cout << "beta0: " << fBeta0 << ", rho: " << fRho << " mm" << endl;
  cout << "------ Parallel Momentum --------" << endl;
  cout << "cospl: " << fCospl << "("<< fdCospl << ")" << endl;
  cout << "BGC: " << fBGC << "("<< fdBGC << ")" << endl;
  cout << "pozpl: " << fPoZpl << "("<< fdPoZpl << ") MeV/c" << endl;

  cout << "\033[1m    -------    The RIB specific information    -------    \033[0m\n";
  cout << "x2: " << fX2Arr[0] << " mm, " << fX2Arr[1] << " mm,";
  cout << " dx2: " << fX2Arr[0] - fX2Arr[1] << " mm" << endl;
} // end member function Print
